clc;
clear;

%% main主脚本
tic;
startTime = cputime;
pde = process_pde_t(3);
vector_h = [1/2 1/4 1/8 1/16];
basis_type = "P1";
theta = 1;
if 0 == theta
    disp(["temporal_type: ", "forward Euler"]);
    error("stop");
elseif 1 == theta
    disp(["temporal_type: ", "backward Euler"]);
elseif 0.5 == theta
    disp(["temporal_type: ", "Crank-Nicolson"]);
end

% process_showbasis(basis_type, "x");

err_Linf = zeros(size(vector_h,2),1);
err_L2 = zeros(size(vector_h,2),1);
err_H1 = zeros(size(vector_h,2),1);

for i = 1:size(vector_h,2)
    h = vector_h(i);
    mesh = process_mesh(pde, h, basis_type);

    % process_showmesh(mesh);

    M = assemble_mass_matrix(mesh);
    A1 = assemble_stiffness_matrix(mesh, basis_type, "dx", basis_type, "dx", pde.cor);
    A2 = assemble_stiffness_matrix(mesh, basis_type, "dy", basis_type, "dy", pde.cor);
    A = A1 + A2;
    U1 = pde.u(mesh.node(:,1),mesh.node(:,2),0);
    if 1 == theta
        if "P1" == basis_type || "P1b" == basis_type
            dt = 4*h^2;
        elseif "P2" == basis_type
            dt = 8*h^3;
        end
    elseif 0.5 == theta
        if "P1" == basis_type || "P1b" == basis_type
            dt = h;
        elseif "P2" == basis_type
            dt = sqrt(h^3);
        end
    end
    Nt = (pde.domain(6) - pde.domain(5))/dt;
    for m = 0:Nt-1
        t1 = m*dt;
        t2 = (m+1)*dt;
        b1 = assemble_load_vector(mesh, pde, t1);
        b2 = assemble_load_vector(mesh, pde, t2);
        [A, b1] = process_boundary_t(A, b1, mesh, pde, t1);
        [A, b2] = process_boundary_t(A, b2, mesh, pde, t2);
        A_hat = M/dt + theta*A;
        b_hat = theta*b2 + (1-theta)*b1 + (M/dt - (1-theta)*A)*U1;
        U2 = A_hat\b_hat;
        U1 = U2;
    end

    process_plot(mesh, pde, U1, basis_type);

    err_Linf(i,1) = process_error_t(mesh, pde, U1, "Linf");
    err_L2(i,1) = process_error_t(mesh, pde, U1, "L2");
    err_H1(i,1) = process_error_t(mesh, pde, U1, "H1");

    fprintf("已完成\t%d/%d\n", i, size(vector_h,2));
end

process_rate(err_Linf, vector_h);
process_rate(err_L2, vector_h);
process_rate(err_H1, vector_h);

elapsedTOCTime = toc;
elapsedCPUTime = cputime - startTime;
disp(["TOC time(s)", num2str(elapsedTOCTime)]);
disp(["CPU time(s)", num2str(elapsedCPUTime)]);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%% 以下是脚本所用到的函数 %%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% process_pde_t
function PDE = process_pde_t(index)
switch index
    case 1
        PDE.index = index;
        PDE.u = @(x,y,t) exp(x+y+t);
        PDE.dudx = @(x,y,t) exp(x+y+t);
        PDE.dudy = @(x,y,t) exp(x+y+t);
        PDE.cor = 2;
        PDE.f = @(x,y,t) -3*exp(x+y+t);
        PDE.domain = [0, 2, 0, 1, 0, 1];
        PDE.bdary = [0, 0, 0, 0];
    case 2
        PDE.index = index;
        PDE.u = @(x,y,t) exp(x+y+t);
        PDE.dudx = @(x,y,t) exp(x+y+t);
        PDE.dudy = @(x,y,t) exp(x+y+t);
        PDE.cor = 2;
        PDE.f = @(x,y,t) -3*exp(x+y+t);
        PDE.domain = [0, 2, 0, 1, 0, 1];
        PDE.bdary = [0, 1, 0, 0];
    case 3
        PDE.index = index;
        PDE.u = @(x,y,t) exp(0.5*(x+y)-t);
        PDE.dudx = @(x,y,t) 0.5*exp(0.5*(x+y)-t);
        PDE.dudy = @(x,y,t) 0.5*exp(0.5*(x+y)-t);
        PDE.cor = 1;
        PDE.f = @(x,y,t) -1.5*exp(0.5*(x+y)-t);
        PDE.domain = [0, 1, 0, 1, 0, 1];
        PDE.bdary = [0, 0, 0, 0];
        % case 5
        %     PDE.index = index;
        %     PDE.u = @(x,y,t) exp(x+y+t);
        %     PDE.dudx = @(x,y,t) exp(x+y+t);
        %     PDE.dudy = @(x,y,t) exp(x+y+t);
        %     PDE.cor = 2;
        %     PDE.f = @(x,y,t) -3*exp(x+y+t);
        %     PDE.r = -1;
        %     PDE.domain = [0, 2, 0, 1, 0, 1];
        %     PDE.bdary = [0, 2, 0, 0];
    otherwise
        error("Invalid PDE index.");
end
fprintf("PDE index: %d\n", PDE.index);
fprintf("PDE domain: xmin,xmax,ymin,ymax,tmix,tmax\n");
disp(PDE.domain);
fprintf("PDE boundary: bottom,right,top,left (0=Dirichlet,1=Neumann,2=Robin)\n");
disp(PDE.bdary);
end

%% assemble_mass_matrix
function matrix_M = assemble_mass_matrix(mesh)
Nb1 = mesh.Nb;
Nlb1 = mesh.Nlb;
Nb2 = Nb1;
Nlb2 = Nlb1;
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = phi(:,j) .* phi(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_M = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_stiffness_matrix
function matrix_A = assemble_stiffness_matrix(mesh, B1_type, D1_type, B2_type, D2_type, cor)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points, gauss_weights] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points(:,1), gauss_points(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points(:,1), gauss_points(:,2), B2_type, D2_type);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = cor * phi1(:,j) .* phi2(:,i);
            r = gauss_weights * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_load_vector
function vector_b = assemble_load_vector(mesh, pde, t)
vector_b = zeros(mesh.Nb,1);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
    ff = pde.f(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    for i = 1:mesh.Nlb
        f = ff .* phi(:,i);
        r = gauss_weights_loc * f;
        vector_b(mesh.elem(n,i),1) = vector_b(mesh.elem(n,i),1) + r;
    end
end
end

%% process_boundary
function [A, b] = process_boundary_t(A, b, mesh, pde, t)
% Robin
% matrix_A = sparse(mesh.Nb, mesh.Nb);
% vector_b = zeros(mesh.Nb,1);
% for k = 1:size(pde.bdary,2)
%     if 2 == pde.bdary(1,k)
%         bde = boundary_edges(mesh, k);
%         rows = zeros(size(bde,1),1);
%         cols = rows;
%         values = rows;
%         flag = 1;
%         for n = 1:size(bde,1)
%             points = mesh.node(bde(n,2:3),:);
%             AB = [points(2,1)-points(1,1); points(2,2)-points(1,2)];
%             vector_n = [0, 1; -1, 0] * AB;
%             vector_n = vector_n./norm(vector_n);
%             fq = @(x,y,t) vector_n(1)*pde.dudx(x,y,t) + vector_n(2)*pde.dudy(x,y,t) + pde.r*pde.u(x,y,t);
%             nk = bde(n,1);
%             E = mesh.node(mesh.elem(nk,:),:);
%             V = E(1:3,:);
%             [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(points, "1d");
%             phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
%             ffq = fq(gauss_points_loc(:,1),gauss_points_loc(:,2),t);
%             for i = 1:mesh.Nlb
%                 f = pde.cor * ffq .* phi(:,i);
%                 r = gauss_weights_loc * f;
%                 vector_b(mesh.elem(nk,i),1) = vector_b(mesh.elem(nk,i),1) + r;
%                 for j = 1:mesh.Nlb
%                     f = pde.cor * pde.r * phi(:,j) .* phi(:,i);
%                     r = gauss_weights_loc * f;
%                     rows(flag) = mesh.elem(nk,i);
%                     cols(flag) = mesh.elem(nk,j);
%                     values(flag) = r;
%                     flag = flag + 1;
%                 end
%             end
%         end
%         matrix_A = sparse(rows, cols, values, mesh.Nb, mesh.Nb);
%     end
% end
% A = A + matrix_A;
% b = b + vector_b;
% Neumann
vector_b = zeros(mesh.Nb,1);
for k = 1:size(pde.bdary,2)
    if 1 == pde.bdary(1,k)
        bde = boundary_edges(mesh, k);
        for n = 1:size(bde,1)
            points = mesh.node(bde(n,2:3),:);
            AB = [points(2,1)-points(1,1); points(2,2)-points(1,2)];
            vector_n = [0, 1; -1, 0] * AB;
            vector_n = vector_n./norm(vector_n);
            fp = @(x,y,t) vector_n(1)*pde.dudx(x,y,t) + vector_n(2)*pde.dudy(x,y,t);
            nk = bde(n,1);
            E = mesh.node(mesh.elem(nk,:),:);
            V = E(1:3,:);
            [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(points, "1d");
            phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
            ffp = fp(gauss_points_loc(:,1),gauss_points_loc(:,2),t);
            for i = 1:mesh.Nlb
                f = pde.cor * ffp .* phi(:,i);
                r = gauss_weights_loc * f;
                vector_b(mesh.elem(nk,i),1) = vector_b(mesh.elem(nk,i),1) + r;
            end
        end
    end
end
b = b + vector_b;
% Dirichlet
for k = 1:size(pde.bdary,2)
    if 0 == pde.bdary(1,k)
        bdn = boundary_nodes(mesh, k);
        A(bdn,:) = 0;
        A(sub2ind(size(A), bdn, bdn)) = 1;
        b(bdn) = pde.u(mesh.node(bdn,1),mesh.node(bdn,2),t);
    end
end
end

%% process_error
function error_value = process_error_t(mesh, pde, U, error_type)
t_end = pde.domain(6);
switch error_type
    case "Linf"
        err_E = 0;
        for n = 1:mesh.Ne
            E = mesh.node(mesh.elem(n,:),:);
            V = E(1:3,:);
            [gauss_points, ~] = gauss_lengdre_loc(V, "2d");
            phi = basis_function_loc(V, gauss_points(:,1), gauss_points(:,2), mesh.basis_type, "x");
            uh_E = phi * U(mesh.elem(n,:));
            f = max(abs(pde.u(gauss_points(:,1),gauss_points(:,2),t_end) - uh_E));
            if f > err_E
                err_E = f;
            end
        end
        error_value = err_E;
    case "L2"
        err_E = 0;
        for n = 1:mesh.Ne
            E = mesh.node(mesh.elem(n,:),:);
            V = E(1:3,:);
            [gauss_points, gauss_weights] = gauss_lengdre_loc(V, "2d");
            phi = basis_function_loc(V, gauss_points(:,1), gauss_points(:,2), mesh.basis_type, "x");
            uh_E = phi * U(mesh.elem(n,:));
            f = (pde.u(gauss_points(:,1),gauss_points(:,2),t_end) - uh_E).^2;
            r = gauss_weights * f;
            err_E = err_E + r;
        end
        error_value = sqrt(err_E);
    case "H1"
        err_Ex = 0;
        err_Ey = 0;
        for n = 1:mesh.Ne
            E = mesh.node(mesh.elem(n,:),:);
            V = E(1:3,:);
            [gauss_points, gauss_weights] = gauss_lengdre_loc(V, "2d");
            phix = basis_function_loc(V, gauss_points(:,1), gauss_points(:,2), mesh.basis_type, "dx");
            phiy = basis_function_loc(V, gauss_points(:,1), gauss_points(:,2), mesh.basis_type, "dy");
            uh_Ex = phix * U(mesh.elem(n,:));
            uh_Ey = phiy * U(mesh.elem(n,:));
            f1 = (pde.dudx(gauss_points(:,1),gauss_points(:,2),t_end) - uh_Ex).^2;
            f2 = (pde.dudy(gauss_points(:,1),gauss_points(:,2),t_end) - uh_Ey).^2;
            rx = gauss_weights * f1;
            ry = gauss_weights * f2;
            err_Ex = err_Ex + rx;
            err_Ey = err_Ey + ry;
        end
        error_value = sqrt(err_Ex + err_Ey);
    otherwise
        error("Invalid error type.");
end
end